import os
import sys
import csv
import time
import json
from core.deploy import MiniCPMVDeploy
from inference.video_inference import VideoInferencer
from utils.batch_processor import process_directory

class BatchVideoInference:
    """
    MiniCPM-V 批量视频推理工具，基于多任务推理服务器实现模型一次加载多次使用
    """
    def __init__(self, config_file=None):
        # 初始化部署类（不立即加载模型）
        self.deploy = MiniCPMVDeploy(config_file)
        self.model_loaded = False
        
        # 从配置中加载视频处理参数
        if config_file and os.path.exists(config_file):
            try:
                with open(config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    self.video_sample_interval = config.get('video_config', {}).get('sample_interval', 60)
                    self.video_max_frames = config.get('video_config', {}).get('max_frames', 8)
            except Exception as e:
                print(f"加载配置文件失败: {str(e)}")
                self.video_sample_interval = 60
                self.video_max_frames = 8
        else:
            self.video_sample_interval = 60
            self.video_max_frames = 8
        
    def load_model(self):
        """加载模型（只需执行一次）"""
        if not self.model_loaded:
            print("\n===== 开始加载模型 =====")
            start_time = time.time()
            if self.deploy.load_model():
                self.model_loaded = True
                load_time = time.time() - start_time
                print(f"模型加载成功，耗时: {load_time:.2f}秒")
                print("模型已加载至内存，可进行多次推理任务")
            else:
                print("模型加载失败，请检查模型配置和路径")
                return False
        else:
            print("模型已经加载，无需重复加载")
        return True
        
    def batch_process(self, folder_path, output_csv, prompt_template=None, frame_interval=10, max_workers=None, recursive=True, no_save=False):
        """
        批量处理视频文件夹，利用process_directory函数实现模型一次加载多次使用
        
        Args:
            folder_path: 视频文件夹路径
            output_csv: 输出CSV文件路径
            prompt_template: 提示词模板，支持变量替换
            frame_interval: 抽帧间隔，默认使用配置中的sample_interval
            max_workers: 最大工作线程数
            recursive: 是否递归搜索子目录
            no_save: 是否不保存单个结果文件
        
        Returns:
            是否处理成功
        """
        if not self.model_loaded:
            print("错误: 模型尚未加载，请先加载模型")
            return False
        
        # 检查文件夹是否存在
        if not os.path.exists(folder_path):
            print(f"错误: 文件夹 {folder_path} 不存在")
            return False
        
        # 如果没有提供提示词，使用默认提示词
        if not prompt_template:
            prompt_template = "请详细描述这个视频的内容，包括主要场景、动作和关键事件"
        
        # 如果没有提供抽帧间隔，使用配置中的值
        if frame_interval is None:
            frame_interval = self.video_sample_interval
        
        print(f"\n===== 开始批量视频推理 =====")
        print(f"视频目录: {folder_path}")
        print(f"提示词模板: {prompt_template}")
        print(f"抽帧间隔: {frame_interval}")
        print(f"是否递归搜索: {recursive}")
        if max_workers is not None:
            print(f"最大工作线程数: {max_workers}")
        if output_csv:
            print(f"CSV导出路径: {output_csv}")
        print(f"是否保存单个结果文件: {not no_save}")
        
        # 开始批量处理，使用process_directory函数
        start_time = time.time()
        results = process_directory(
            dir_path=folder_path,
            inferencer_class=VideoInferencer,
            prompt=prompt_template,
            save_result=not no_save,
            export_csv=output_csv,
            max_workers=max_workers,
            recursive=recursive,
            deploy=self.deploy,  # 传递已加载的deploy实例
            frame_interval=frame_interval,  # 视频处理特有参数
            video_max_frames=self.video_max_frames  # 最大帧数参数
        )
        
        infer_time = time.time() - start_time
        print(f"\n批量视频推理完成，耗时: {infer_time:.2f}秒")
        
        # 统计信息
        if results:
            total_files = len(results)
            success_count = sum(1 for _, result, _ in results if result and "处理失败" not in str(result))
            failed_count = total_files - success_count
            total_time_cost = sum(time_cost for _, _, time_cost in results)
            
            print(f"统计信息:")
            print(f"- 总视频数: {total_files}")
            print(f"- 成功处理: {success_count}")
            print(f"- 失败处理: {failed_count}")
            print(f"- 总耗时: {total_time_cost:.2f}秒")
            print(f"- 平均耗时: {total_time_cost/total_files:.2f}秒/个")
        
        return True
        
    def unload_model(self):
        """卸载模型，释放内存"""
        if self.model_loaded:
            self.deploy.unload_model()
            self.model_loaded = False
            print("模型已卸载，内存已释放")
        else:
            print("模型尚未加载，无需卸载")

def main():
    """主函数"""
    print("===== MiniCPM-V 批量视频推理工具 =====")
    print("基于多任务推理服务器实现模型一次加载，多次使用")
    
    # 创建批量推理实例
    batch_infer = BatchVideoInference()
    
    # 加载模型（只需加载一次）
    print("\n正在加载模型...")
    if not batch_infer.load_model():
        print("模型加载失败，程序退出")
        return
    
    # 获取用户输入
    print("\n请输入视频文件夹路径: ")
    folder_path = input().strip()
    if not folder_path:
        print("错误: 视频文件夹路径不能为空")
        return
    
    # 验证文件夹路径
    if not os.path.isdir(folder_path):
        print(f"错误: {folder_path} 不是有效的文件夹路径")
        return
    
    # 获取输出CSV文件路径
    print("请输入输出CSV文件路径 (默认: outputs/batch_video_results.csv): ")
    output_csv = input().strip()
    if not output_csv:
        output_csv = "outputs/batch_video_results.csv"
    
    # 获取提示词模板
    print("请输入推理提示词模板 (可选，支持变量替换): ")
    prompt_template = input().strip()
    
    # 获取并行线程数
    print("请输入并行处理线程数 (可选，留空自动调整): ")
    max_workers_input = input().strip()
    max_workers = None
    if max_workers_input and max_workers_input.isdigit():
        max_workers = int(max_workers_input)
    
    # 获取抽帧间隔
    print("请输入视频抽帧间隔 (默认: 60，即每60帧抽1帧): ")
    frame_interval_input = input().strip()
    frame_interval = 60  # 默认值
    if frame_interval_input and frame_interval_input.isdigit():
        frame_interval = int(frame_interval_input)
    
    # 是否递归搜索子目录
    print("是否递归搜索子目录？(y/n，默认: y): ")
    recursive_input = input().strip().lower()
    recursive = True
    if recursive_input == 'n':
        recursive = False
    
    # 是否保存单个结果文件
    print("是否保存单个结果文件？(y/n，默认: y): ")
    save_single_input = input().strip().lower()
    no_save = False
    if save_single_input == 'n':
        no_save = True
    
    # 开始批量处理
    print("\n开始批量处理视频...")
    start_total_time = time.time()
    
    if batch_infer.batch_process(
        folder_path=folder_path,
        output_csv=output_csv,
        prompt_template=prompt_template,
        frame_interval=frame_interval,
        max_workers=max_workers,
        recursive=recursive,
        no_save=no_save
    ):
        end_total_time = time.time()
        print(f"\n批量处理完成！总耗时: {end_total_time - start_total_time:.2f}秒")
        
        # 询问是否继续处理其他视频或卸载模型
        while True:
            print("\n请选择操作:")
            print("1. 处理其他视频文件夹")
            print("2. 卸载模型并退出")
            choice = input().strip()
            
            if choice == '1':
                # 处理其他视频文件夹
                print("请输入新的视频文件夹路径: ")
                new_folder_path = input().strip()
                if not new_folder_path:
                    print("错误: 视频文件夹路径不能为空")
                    continue
                if not os.path.isdir(new_folder_path):
                    print(f"错误: {new_folder_path} 不是有效的文件夹路径")
                    continue
                
                print("请输入新的输出CSV文件路径: ")
                new_output_csv = input().strip()
                if not new_output_csv:
                    # 生成默认文件名
                    import datetime
                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                    new_output_csv = f"outputs/batch_video_results_{timestamp}.csv"
                
                print("\n开始处理新的视频文件夹...")
                batch_infer.batch_process(
                    folder_path=new_folder_path,
                    output_csv=new_output_csv,
                    prompt_template=prompt_template,
                    frame_interval=frame_interval,
                    max_workers=max_workers,
                    recursive=recursive,
                    no_save=no_save
                )
            elif choice == '2':
                # 卸载模型并退出
                batch_infer.unload_model()
                print("程序已退出")
                break
            else:
                print("无效的选择，请重新输入")
    else:
        print("批量处理失败")
        batch_infer.unload_model()

if __name__ == "__main__":
    main()